#    Copyright (c) 2010 Abhishek Patnia (patnia@isi.edu) and Tarang Desai (tarangde@usc.edu)
#
#    Permission is hereby granted, free of charge, to any person obtaining a copy
#    of this software and associated documentation files (the "Software"), to deal
#    in the Software without restriction, including without limitation the rights
#    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#    copies of the Software, and to permit persons to whom the Software is
#    furnished to do so, subject to the following conditions:
#
#    The above copyright notice and this permission notice shall be included in
#    all copies or substantial portions of the Software.
#
#    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
#    THE SOFTWARE.

from __future__ import division
from collections import defaultdict
import gflags as flags
import sys

FLAGS=flags.FLAGS

# global constants
GIVEN = '->'
NULL = 'NULL'

def createCorpus(eng_fn, hindi_fn):
    """
        Read test data and create a test corpus
    """
    # create a dict of corpus with key as sentence number and values as English Hindi sentence pair (tuple)
    corpus = {}
    for line in open(eng_fn):
        num = int(line.strip().split('>')[0].split('=')[1]) # sentence number
        eng_sent = [NULL]
        for eng_word in line.strip().split('>')[1].split('<')[0].split(' '):
            eng_word = eng_word.strip()
            if len(eng_word) > 0:
                eng_sent.append(eng_word)
        corpus[num] = eng_sent

    for line in open(hindi_fn):
        num = int(line.strip().split('>')[0].split('=')[1]) # sentence number
        hind_sent = []
        for hind_word in line.strip().split('>')[1].split('<')[0].split(' '):
            hind_word = hind_word.strip()
            if len(hind_word) > 0:
                hind_sent.append(hind_word)
        corpus[num] = (corpus[num], hind_sent)

    return corpus

def readProbability(probs_fn):
    """
        Read probabilities from the given file. Format is key=value
    """
    probDict = defaultdict(float)
    for line in open(probs_fn):
        key, prob = line.strip().split('=')
        probDict[key] = float(prob)
    return probDict

def estimateBestAlignment(hindi, english, T):
    """
        Estimate the best alignment of the hindi and english sentence pair for the given transittion probabilities.
    """
    alignment = []
    for hindi_word in hindi:
        max_alignment = (-1, -1)
        for i, eng_word in enumerate(english):
            if max_alignment[1] < T[hindi_word + GIVEN + eng_word]:
                max_alignment = (i, T[hindi_word + GIVEN + eng_word])
        alignment.append(max_alignment[0])

    return alignment

def bestAlignment(corpus, T):
    """
        Find the best alignments for the given corpus and transition probabilities and write the results to a file
    """
    wa_file = open(FLAGS.output, 'w')
    for num, item in corpus.items():
        alignment = estimateBestAlignment(item[1], item[0], T)
        engBoolMap = []
        for i in range(len(item[0])):
            engBoolMap.append(False)
        for h, e in enumerate(alignment):
            engBoolMap[e] = True
            wa_file.write('%s %s %s\n' % (str(num), str(e), str(h + 1)))
        for i, item in enumerate(engBoolMap):
            if not(item):
                wa_file.write('%s %s %s\n' % (str(num), str(i), 0))
    wa_file.flush()
    wa_file.close()

if __name__ == '__main__':

    # Flags
    # English data file
    flags.DEFINE_string('engTestFN', 'data/test/test.e', 'English test data filename')
    # Hindi data file
    flags.DEFINE_string('hindTestFN', 'data/test/test.h', 'Hindi test data filename')
    # Input probabilities filename
    flags.DEFINE_string('input', 'models/transition.model1.prob', 'Model 1 transition probabilities')
    # Output filename
    flags.DEFINE_string('output', 'results/align.ibm.model1.test.wa', 'Final best alignments for the test sentences')

    # initialize flags from command line
    FLAGS(sys.argv)

    corpus = createCorpus(FLAGS.engTestFN, FLAGS.hindTestFN)
    T = readProbability(FLAGS.input)
    bestAlignment(corpus, T)